from sklearn.linear_model import LinearRegression, Ridge, BayesianRidge
import pandas as pd
from .classification import split_dataset, get_X_Y_data
from .visualization import get_heatmap_path, get_feature_imp_path
from hte.error.handle import abort_on_error
from hte.error.models import HTEError


def linear_regression(params):
    clf = LinearRegression()
    clf.set_params(**params)
    return clf


def linear_result(task, clf, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    if Y is None:
        abort_on_error(HTEError.BAD_DATA)
    x_train, x_test, y_train, y_test = split_dataset(X, Y, 0.3)
    clf.fit(x_train, y_train)
    result = {}
    result['score'] = clf.score(x_test, y_test)
    result['coef'] = list(clf.coef_)
    result['intercept'] = clf.intercept_
    imgs = []
    imgs.append(get_heatmap_path(X))
    return result, imgs


def linear_imgs(task, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    imgs = []
    imgs.append(get_heatmap_path(X))
    return imgs


def ridg_regression(params):
    clf = Ridge()
    clf.set_params(**params)
    return clf


def ridge_result(task, clf, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    if Y is None:
        abort_on_error(HTEError.BAD_DATA)
    x_train, x_test, y_train, y_test = split_dataset(X, Y, 0.3)
    clf.fit(x_train, y_train)
    result = {}
    result['score'] = clf.score(x_test, y_test)
    result['coef'] = list(clf.coef_)
    result['intercept'] = clf.intercept_
    imgs = []
    imgs.append(get_heatmap_path(X))
    return result, imgs


def ridge_imgs(task, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    imgs = []
    imgs.append(get_heatmap_path(X))
    return imgs


def bayes_regression(params):
    clf = BayesianRidge()
    clf.set_params(**params)
    return clf


def bayes_result(task, clf, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    if Y is None:
        abort_on_error(HTEError.BAD_DATA)
    x_train, x_test, y_train, y_test = split_dataset(X, Y, 0.3)
    clf.fit(x_train, y_train)
    result = {}
    result['score'] = clf.score(x_test, y_test)
    result['coef'] = list(clf.coef_)
    result['intercept'] = clf.intercept_
    imgs = []
    imgs.append(get_heatmap_path(X))
    return result, imgs


def bayes_imgs(task, dataset):
    X, Y = get_X_Y_data(dataset, task.label, task.features)
    imgs = []
    imgs.append(get_heatmap_path(X))
    return imgs


if __name__ == "__main__":
    df = pd.read_excel('linear.xlsx')
    # df.hist()
    # indexs = df.columns
    # y_label = 'Y'
    # x_label = indexs.drop([y_label])
    # X_data = df[x_label]
    # Y_data = df[y_label]
    # x_train, x_test, y_train, y_test = train_test_split(
    #     X_data, Y_data, test_size=0.3, random_state=618)
    # clf = linear_regression({})
    # clf.fit(x_train, y_train)
    # score = clf.score(x_test, y_test)
    # result = clf.predict(x_test)
    # plt.figure()
    # plt.plot(x_test, y_test, 'go-', label='true value')
    # plt.plot(x_test, result, 'ro-', label='predict value')
    # plt.title('score: %f' % score)
    # plt.legend()
    # plt.show()
    clf = linear_regression({})
    X, Y = get_X_Y_data(df, ['Y'])
    x_train, x_test, y_train, y_test = split_dataset(X, Y, 0.3)
    clf.fit(x_train, y_train)
    result = {}
    result['score'] = clf.score(x_test, y_test)
    result['coef'] = clf.coef_
    result['intercept'] = clf.intercept_
    imgs = []
    imgs.append(get_heatmap_path(X))
    imgs.append(get_feature_imp_path(X, Y))
    print(result)
    print(imgs)
